#!/usr/bin/env python3
"""
MCP 基础数学运算服务器
提供基本的数学运算功能：加减乘除乘方、表达式计算、线性代数运算等
使用stdio传输协议与客户端通信
"""

from mcp.server.fastmcp import FastMCP
import math
import numpy as np
from typing import Union, List, Dict, Any

# 创建MCP服务器实例
mcp = FastMCP("mcp-basic-math")

@mcp.tool()
def add(a: float, b: float) -> float:
    """
    执行加法运算
    
    Args:
        a: 第一个数字
        b: 第二个数字
    
    Returns:
        两个数字的和
    """
    result = a + b
    return result

@mcp.tool()
def subtract(a: float, b: float) -> float:
    """
    执行减法运算
    
    Args:
        a: 被减数
        b: 减数
    
    Returns:
        两个数字的差
    """
    result = a - b
    return result

@mcp.tool()
def multiply(a: float, b: float) -> float:
    """
    执行乘法运算
    
    Args:
        a: 第一个数字
        b: 第二个数字
    
    Returns:
        两个数字的积
    """
    result = a * b
    return result

@mcp.tool()
def divide(a: float, b: float) -> Union[float, str]:
    """
    执行除法运算
    
    Args:
        a: 被除数
        b: 除数
    
    Returns:
        两个数字的商，如果除数为0则返回错误信息
    """
    if b == 0:
        return "错误：除数不能为零"
    result = a / b
    return result

@mcp.tool()
def power(base: float, exponent: float) -> float:
    """
    执行乘方运算
    
    Args:
        base: 底数
        exponent: 指数
    
    Returns:
        底数的指数次幂
    """
    result = math.pow(base, exponent)
    return result

@mcp.tool()
def square_root(number: float) -> Union[float, str]:
    """
    计算平方根
    
    Args:
        number: 要计算平方根的数字
    
    Returns:
        数字的平方根，如果数字为负数则返回错误信息
    """
    if number < 0:
        return "错误：不能计算负数的平方根"
    result = math.sqrt(number)
    return result

@mcp.tool()
def get_server_info() -> dict:
    """
    获取MCP服务器的基本信息
    
    Returns:
        包含服务器信息的字典
    """
    return {
        "server_name": "mcp-basic-math",
        "version": "1.0.0",
        "description": "MCP基础数学运算服务器，提供基础数学运算和线性代数计算功能",
        "transport": "stdio",
        "modules": [
            "基础数学运算模块",
            "线性代数计算模块"
        ],
        "available_tools": [
            # 基础运算
            "add", "subtract", "multiply", "divide", "power", "square_root", "calculate_expression",
            # 线性代数
            "gaussian_elimination", "lu_decomposition", "qr_decomposition", "power_method", "jacobi_method",
            "matrix_operations", "eigenvalue_analysis", "svd_decomposition",
            "get_server_info"
        ]
    }

@mcp.tool()
def gaussian_elimination(matrix: list, vector: list) -> Union[list, str]:
    """
    使用高斯消元法求解线性方程组 Ax = b
    
    Args:
        matrix: 系数矩阵 A (二维列表)
        vector: 常数向量 b (一维列表)
    
    Returns:
        方程组的解向量，如果无解或有无穷解则返回错误信息
    """
    try:
        import copy
        
        # 验证输入
        if not matrix or not vector:
            return "错误：矩阵或向量为空"
        
        n = len(matrix)
        if len(vector) != n:
            return "错误：矩阵行数与向量长度不匹配"
        
        for row in matrix:
            if len(row) != n:
                return "错误：矩阵不是方阵"
        
        # 创建增广矩阵
        augmented = []
        for i in range(n):
            row = copy.deepcopy(matrix[i])
            row.append(float(vector[i]))
            augmented.append(row)
        
        # 前向消元
        for i in range(n):
            # 寻找主元
            max_row = i
            for k in range(i + 1, n):
                if abs(augmented[k][i]) > abs(augmented[max_row][i]):
                    max_row = k
            
            # 交换行
            if max_row != i:
                augmented[i], augmented[max_row] = augmented[max_row], augmented[i]
            
            # 检查是否为奇异矩阵
            if abs(augmented[i][i]) < 1e-10:
                return "错误：矩阵奇异，无唯一解"
            
            # 消元
            for k in range(i + 1, n):
                factor = augmented[k][i] / augmented[i][i]
                for j in range(i, n + 1):
                    augmented[k][j] -= factor * augmented[i][j]
        
        # 回代求解
        solution = [0.0] * n
        for i in range(n - 1, -1, -1):
            solution[i] = augmented[i][n]
            for j in range(i + 1, n):
                solution[i] -= augmented[i][j] * solution[j]
            solution[i] /= augmented[i][i]
        
        return solution
        
    except Exception as e:
        return f"错误：高斯消元法计算失败 - {str(e)}"

@mcp.tool()
def lu_decomposition(matrix: list) -> Union[dict, str]:
    """
    LU分解：将矩阵分解为下三角矩阵L和上三角矩阵U
    
    Args:
        matrix: 要分解的方阵 (二维列表)
    
    Returns:
        包含L矩阵和U矩阵的字典，如果分解失败则返回错误信息
    """
    try:
        import copy
        
        if not matrix:
            return "错误：矩阵为空"
        
        n = len(matrix)
        for row in matrix:
            if len(row) != n:
                return "错误：矩阵不是方阵"
        
        # 初始化L和U矩阵
        L = [[0.0] * n for _ in range(n)]
        U = [[0.0] * n for _ in range(n)]
        
        # LU分解
        for i in range(n):
            # 计算U矩阵的第i行
            for j in range(i, n):
                U[i][j] = float(matrix[i][j])
                for k in range(i):
                    U[i][j] -= L[i][k] * U[k][j]
            
            # 计算L矩阵的第i列
            L[i][i] = 1.0  # L矩阵对角线元素为1
            for j in range(i + 1, n):
                if abs(U[i][i]) < 1e-10:
                    return "错误：矩阵奇异，无法进行LU分解"
                L[j][i] = float(matrix[j][i])
                for k in range(i):
                    L[j][i] -= L[j][k] * U[k][i]
                L[j][i] /= U[i][i]
        
        return {
            "L_matrix": L,
            "U_matrix": U
        }
        
    except Exception as e:
        return f"错误：LU分解失败 - {str(e)}"

@mcp.tool()
def qr_decomposition(matrix: list) -> Union[dict, str]:
    """
    QR分解：将矩阵分解为正交矩阵Q和上三角矩阵R
    
    Args:
        matrix: 要分解的矩阵 (二维列表)
    
    Returns:
        包含Q矩阵和R矩阵的字典，如果分解失败则返回错误信息
    """
    try:
        import copy
        import math
        
        if not matrix:
            return "错误：矩阵为空"
        
        m = len(matrix)
        n = len(matrix[0])
        
        # 转置矩阵以便按列处理
        A = [[float(matrix[i][j]) for i in range(m)] for j in range(n)]
        Q = [[0.0] * m for _ in range(m)]
        R = [[0.0] * n for _ in range(n)]
        
        # Gram-Schmidt正交化
        for j in range(n):
            # 复制当前列
            v = A[j][:]
            
            # 减去之前所有正交向量的投影
            for i in range(j):
                # 计算投影系数
                dot_product = sum(Q[i][k] * A[j][k] for k in range(m))
                R[i][j] = dot_product
                
                # 减去投影
                for k in range(m):
                    v[k] -= dot_product * Q[i][k]
            
            # 计算向量的模长
            norm = math.sqrt(sum(x * x for x in v))
            if norm < 1e-10:
                return "错误：矩阵列向量线性相关，无法进行QR分解"
            
            R[j][j] = norm
            
            # 标准化得到正交向量
            for k in range(m):
                Q[j][k] = v[k] / norm
        
        # 转置Q矩阵
        Q_transposed = [[Q[j][i] for j in range(n)] for i in range(m)]
        
        return {
            "Q_matrix": Q_transposed,
            "R_matrix": R
        }
        
    except Exception as e:
        return f"错误：QR分解失败 - {str(e)}"

@mcp.tool()
def power_method(matrix: list, max_iterations: int = 100, tolerance: float = 1e-6) -> Union[dict, str]:
    """
    使用幂法计算矩阵的最大特征值和对应的特征向量
    
    Args:
        matrix: 输入矩阵 (二维列表)
        max_iterations: 最大迭代次数
        tolerance: 收敛容差
    
    Returns:
        包含最大特征值和特征向量的字典，如果计算失败则返回错误信息
    """
    try:
        import random
        import math
        
        if not matrix:
            return "错误：矩阵为空"
        
        n = len(matrix)
        for row in matrix:
            if len(row) != n:
                return "错误：矩阵不是方阵"
        
        # 初始化随机向量
        x = [random.random() for _ in range(n)]
        
        # 标准化
        norm = math.sqrt(sum(xi * xi for xi in x))
        x = [xi / norm for xi in x]
        
        eigenvalue = 0
        
        for iteration in range(max_iterations):
            # 矩阵向量乘法 y = A * x
            y = [0.0] * n
            for i in range(n):
                for j in range(n):
                    y[i] += float(matrix[i][j]) * x[j]
            
            # 计算瑞利商（特征值的近似）
            new_eigenvalue = sum(x[i] * y[i] for i in range(n))
            
            # 标准化
            norm = math.sqrt(sum(yi * yi for yi in y))
            if norm < 1e-10:
                return "错误：向量收敛到零向量"
            
            x = [yi / norm for yi in y]
            
            # 检查收敛
            if abs(new_eigenvalue - eigenvalue) < tolerance:
                return {
                    "eigenvalue": new_eigenvalue,
                    "eigenvector": x,
                    "iterations": iteration + 1
                }
            
            eigenvalue = new_eigenvalue
        
        return {
            "eigenvalue": eigenvalue,
            "eigenvector": x,
            "iterations": max_iterations,
            "converged": False
        }
        
    except Exception as e:
        return f"错误：幂法计算失败 - {str(e)}"

@mcp.tool()
def jacobi_method(matrix: list, max_iterations: int = 100, tolerance: float = 1e-6) -> Union[dict, str]:
    """
    使用雅可比方法计算对称矩阵的所有特征值和特征向量
    
    Args:
        matrix: 输入的对称矩阵 (二维列表)
        max_iterations: 最大迭代次数
        tolerance: 收敛容差
    
    Returns:
        包含所有特征值和特征向量的字典，如果计算失败则返回错误信息
    """
    try:
        import copy
        import math
        
        if not matrix:
            return "错误：矩阵为空"
        
        n = len(matrix)
        for row in matrix:
            if len(row) != n:
                return "错误：矩阵不是方阵"
        
        # 检查是否为对称矩阵
        for i in range(n):
            for j in range(n):
                if abs(float(matrix[i][j]) - float(matrix[j][i])) > 1e-10:
                    return "错误：矩阵不是对称矩阵"
        
        # 初始化
        A = [[float(matrix[i][j]) for j in range(n)] for i in range(n)]
        V = [[1.0 if i == j else 0.0 for j in range(n)] for i in range(n)]  # 特征向量矩阵
        
        for iteration in range(max_iterations):
            # 找到最大的非对角元素
            max_val = 0
            p, q = 0, 1
            for i in range(n):
                for j in range(i + 1, n):
                    if abs(A[i][j]) > max_val:
                        max_val = abs(A[i][j])
                        p, q = i, j
            
            # 检查收敛
            if max_val < tolerance:
                # 提取特征值（对角元素）
                eigenvalues = [A[i][i] for i in range(n)]
                
                # 提取特征向量（V的列）
                eigenvectors = [[V[i][j] for i in range(n)] for j in range(n)]
                
                return {
                    "eigenvalues": eigenvalues,
                    "eigenvectors": eigenvectors,
                    "iterations": iteration
                }
            
            # 计算旋转角度
            if abs(A[p][p] - A[q][q]) < 1e-10:
                theta = math.pi / 4
            else:
                theta = 0.5 * math.atan(2 * A[p][q] / (A[p][p] - A[q][q]))
            
            c = math.cos(theta)
            s = math.sin(theta)
            
            # 应用Givens旋转
            # 更新矩阵A
            A_pp = A[p][p]
            A_qq = A[q][q]
            A_pq = A[p][q]
            
            A[p][p] = c * c * A_pp + s * s * A_qq - 2 * s * c * A_pq
            A[q][q] = s * s * A_pp + c * c * A_qq + 2 * s * c * A_pq
            A[p][q] = A[q][p] = 0
            
            for i in range(n):
                if i != p and i != q:
                    A_ip = A[i][p]
                    A_iq = A[i][q]
                    A[i][p] = A[p][i] = c * A_ip - s * A_iq
                    A[i][q] = A[q][i] = s * A_ip + c * A_iq
            
            # 更新特征向量矩阵V
            for i in range(n):
                V_ip = V[i][p]
                V_iq = V[i][q]
                V[i][p] = c * V_ip - s * V_iq
                V[i][q] = s * V_ip + c * V_iq
        
        # 如果没有收敛
        eigenvalues = [A[i][i] for i in range(n)]
        eigenvectors = [[V[i][j] for i in range(n)] for j in range(n)]
        
        return {
            "eigenvalues": eigenvalues,
            "eigenvectors": eigenvectors,
            "iterations": max_iterations,
            "converged": False
        }
        
    except Exception as e:
        return f"错误：雅可比方法计算失败 - {str(e)}"

@mcp.tool()
def calculate_expression(expression: str) -> Union[float, str]:
    """
    计算数学表达式
    
    Args:
        expression: 要计算的数学表达式字符串（支持 +, -, *, /, **, (), sqrt, sin, cos, tan, log, exp等数学函数）
    
    Returns:
        表达式的计算结果，如果表达式无效则返回错误信息
    """
    try:
        # 安全的表达式计算，允许数学运算和常用数学函数
        import re
        
        # 检查是否只包含安全的字符和函数
        # 允许：数字、运算符、括号、空格、常用数学函数名
        allowed_pattern = r'^[0-9+\-*/.()\s]*$|sqrt|sin|cos|tan|log|exp|abs|pow|pi|e'
        
        # 预处理：替换一些常见的数学函数和常数
        safe_expression = expression.strip()
        
        # 检查是否包含危险的关键字
        dangerous_keywords = ['import', 'exec', 'eval', '__', 'open', 'file', 'input', 'raw_input']
        if any(keyword in safe_expression.lower() for keyword in dangerous_keywords):
            return "错误：表达式包含不允许的关键字"
        
        # 创建安全的数学环境
        safe_dict = {
            "__builtins__": {},
            "sqrt": math.sqrt,
            "sin": math.sin,
            "cos": math.cos,
            "tan": math.tan,
            "log": math.log,
            "log10": math.log10,
            "exp": math.exp,
            "abs": abs,
            "pow": pow,
            "pi": math.pi,
            "e": math.e,
            "ceil": math.ceil,
            "floor": math.floor,
            "round": round
        }
        
        # 使用 eval 计算表达式（在受控环境中）
        result = eval(safe_expression, safe_dict, {})
        
        if isinstance(result, (int, float)):
            return float(result)
        else:
            return "错误：计算结果不是数字"
            
    except ZeroDivisionError:
        return "错误：除数不能为零"
    except NameError as e:
        return f"错误：未知的函数或变量 - {str(e)}"
    except Exception as e:
        return f"错误：无效的表达式 - {str(e)}"

@mcp.tool()
def matrix_operations(matrix_a: List[List[float]], matrix_b: List[List[float]] = None, operation: str = "add") -> Union[Dict[str, Any], str]:
    """
    执行矩阵运算（加法、减法、乘法、转置、求逆等）
    
    Args:
        matrix_a: 第一个矩阵
        matrix_b: 第二个矩阵（某些运算需要）
        operation: 运算类型 (add, subtract, multiply, transpose, inverse, determinant)
    
    Returns:
        运算结果矩阵或标量值
    """
    try:
        A = np.array(matrix_a)
        
        if operation == "transpose":
            return {"result": A.T.tolist(), "operation": "矩阵转置"}
        
        elif operation == "determinant":
            if A.shape[0] != A.shape[1]:
                return "错误：只能计算方阵的行列式"
            det = np.linalg.det(A)
            return {"result": float(det), "operation": "行列式计算"}
        
        elif operation == "inverse":
            if A.shape[0] != A.shape[1]:
                return "错误：只能计算方阵的逆矩阵"
            try:
                inv_A = np.linalg.inv(A)
                return {"result": inv_A.tolist(), "operation": "矩阵求逆"}
            except np.linalg.LinAlgError:
                return "错误：矩阵奇异，不可逆"
        
        elif matrix_b is not None:
            B = np.array(matrix_b)
            
            if operation == "add":
                if A.shape != B.shape:
                    return "错误：矩阵维度不匹配，无法相加"
                return {"result": (A + B).tolist(), "operation": "矩阵加法"}
            
            elif operation == "subtract":
                if A.shape != B.shape:
                    return "错误：矩阵维度不匹配，无法相减"
                return {"result": (A - B).tolist(), "operation": "矩阵减法"}
            
            elif operation == "multiply":
                if A.shape[1] != B.shape[0]:
                    return "错误：矩阵维度不匹配，无法相乘"
                return {"result": (A @ B).tolist(), "operation": "矩阵乘法"}
        
        return "错误：不支持的运算类型或缺少必要参数"
        
    except Exception as e:
        return f"错误：矩阵运算失败 - {str(e)}"

@mcp.tool()
def svd_decomposition(matrix: List[List[float]]) -> Union[Dict[str, Any], str]:
    """
    奇异值分解（SVD）
    
    Args:
        matrix: 输入矩阵
    
    Returns:
        包含U、S、V矩阵的字典
    """
    try:
        A = np.array(matrix)
        U, S, Vt = np.linalg.svd(A)
        
        return {
            "U_matrix": U.tolist(),
            "singular_values": S.tolist(),
            "V_transpose": Vt.tolist(),
            "rank": np.linalg.matrix_rank(A),
            "condition_number": float(np.linalg.cond(A))
        }
        
    except Exception as e:
        return f"错误：SVD分解失败 - {str(e)}"

@mcp.tool()
def eigenvalue_analysis(matrix: List[List[float]]) -> Union[Dict[str, Any], str]:
    """
    特征值和特征向量分析
    
    Args:
        matrix: 输入方阵
    
    Returns:
        特征值、特征向量及相关分析
    """
    try:
        A = np.array(matrix)
        
        if A.shape[0] != A.shape[1]:
            return "错误：只能计算方阵的特征值"
        
        eigenvalues, eigenvectors = np.linalg.eig(A)
        
        # 按特征值大小排序
        idx = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]
        
        return {
            "eigenvalues": eigenvalues.real.tolist(),
            "eigenvectors": eigenvectors.real.tolist(),
            "trace": float(np.trace(A)),
            "determinant": float(np.linalg.det(A)),
            "is_symmetric": bool(np.allclose(A, A.T)),
            "is_positive_definite": bool(np.all(eigenvalues.real > 0))
        }
        
    except Exception as e:
        return f"错误：特征值分析失败 - {str(e)}"


if __name__ == "__main__":
    # 启动MCP服务器，使用stdio传输协议
    mcp.run(transport="stdio")